# encoding=utf-8

import cv2 as cv
import numpy as np
from RefAlgorithm.OTSU import Otsu


class FractalDimension:
    def predict(self, img: np.ndarray, th_0: int = -1, REG: float = 1e-1):
        """
        :param img: width, height, channel
        :param th_0: binary threshold, we calculate the loc
                   which less than th。if th is two small,
                   algorithm will send back a 'nan'.
        :param REG: 正则向，大津法用，防止除0
        :return: float, nan
        """
        img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        if th_0 != -1:
            th = th_0
        else:
            th_1, var = Otsu(img, REG)
            th = th_1
        return self.fractal_dimension(img, th)

    def fractal_dimension(self, Z: np.ndarray, threshold):
        def boxcount(Z, k):
            S = np.add.reduceat(
                np.add.reduceat(Z, np.arange(0, Z.shape[0], k), axis=0),
                np.arange(0, Z.shape[1], k), axis=1)

            # We count non-empty (0) and non-full boxes (k*k)
            return len(np.where((S > 0) & (S < k * k))[0])

        # Transform Z into a binary array
        Z = (Z < threshold)

        # Minimal dimension of image
        p = min(Z.shape)

        # Greatest power of 2 less than or equal to p
        n = 2 ** np.floor(np.log(p) / np.log(2))

        # Extract the exponent
        n = int(np.log(n) / np.log(2))

        # Build successive box sizes (from 2**n down to 2**1)
        sizes = 2 ** np.arange(n, 1, -1)

        # Actual box counting with decreasing size
        counts = []
        for size in sizes:
            counts.append(boxcount(Z, size))

        # Fit the successive log(sizes) with log (counts)
        coeffs = np.polyfit(np.log(sizes), np.log(counts), 1)
        return -coeffs[0]
